﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.Threading.Tasks;
using TorchSharp;
using static TorchSharp.torch;

namespace NgpSharp.RayMaching
{
    public class SphFromRayFunction : torch.autograd.SingleTensorFunction<SphFromRayFunction>
    {
        public override string Name => nameof(SphFromRayFunction);

        public override torch.Tensor forward(torch.autograd.AutogradContext ctx, params object[] vars)
        {
            using var c = torch.NewDisposeScope();
            Tensor rays_o = (Tensor)vars[0];
            Tensor rays_d = (Tensor)vars[1];
            float radius = (float)vars[2];

            if (!rays_o.is_cuda) rays_o = rays_o.cuda();
            if (!rays_d.is_cuda) rays_d = rays_d.cuda();

            rays_o = rays_o.contiguous().view(-1, 3);
            rays_d = rays_d.contiguous().view(-1, 3);

            long N = rays_o.shape[0];

            Tensor coords = torch.empty(N, 2, dtype: rays_o.dtype, device: rays_o.device);
            RayMachingLoader.sph_from_ray(rays_o.Handle, rays_d.Handle, radius, (uint)N, coords.Handle);

            return coords.MoveToOuterDisposeScope();
        }

        public override List<torch.Tensor> backward(torch.autograd.AutogradContext ctx, torch.Tensor grad_output)
        {
            grad_output.Dispose();
            return [null, null, null];
        }
    }
}
